package com.culture.bootdemo.controller.解决方案.限流;

import redis.clients.jedis.Jedis;

/**
 * @Author: wen
 * @DateTime: 2024/11/19 13:58
 * @Description: 滑动窗口-实现限流  有序集合 (Sorted Set)
 **/
public class SlidingWindowRateLimiterDemo {
    private final Jedis jedis;

    public SlidingWindowRateLimiterDemo(Jedis jedis) {
        this.jedis = jedis;
    }

    /**
     * @param key            被限制的接口或者组合
     * @param limit          限制次数
     * @param windowInMillis 限制单位时间（毫秒）
     * @return
     */
    public boolean isAllowed(String key, int limit, long windowInMillis) {
        long now = System.currentTimeMillis();
        long windowStart = now - windowInMillis;

        // Redis Sorted Set key
        String redisKey = "rate_limiter:" + key;

        // 移除窗口外的请求
        jedis.zremrangeByScore(redisKey, 0, windowStart);

        // 获取当前窗口内的请求数
        long count = jedis.zcard(redisKey);

        if (count < limit) {
            // 添加当前请求时间戳到窗口
            jedis.zadd(redisKey, now, String.valueOf(now));
            // 设置过期时间
            jedis.expire(redisKey, (int) (windowInMillis / 1000));
            return true;
        } else {
            return false;
        }
    }

    public void runTest() {
        // 连接 Redis
        Jedis jedis = new Jedis("localhost", 6379);
        SlidingWindowRateLimiterDemo rateLimiter = new SlidingWindowRateLimiterDemo(jedis);

        String apiKey = "user:login";
        int limit = 5; // 每秒最多 5 次
        long window = 1000; // 时间窗口 1000 毫秒(1s)

        for (int i = 0; i < 10; i++) {
            if (rateLimiter.isAllowed(apiKey, limit, window)) {
                System.out.println("Request allowed: " + i);
            } else {
                System.out.println("Request denied: " + i);
            }
        }

        jedis.close();
    }
}
